/*
* coder: dzlua
* email: 505544956@qq.com
* time : 2018/05/14 10:28
*/
#ifndef __ICON_STREAMS_HPP__
#define __ICON_STREAMS_HPP__
#pragma once


#include <vector>
#include <memory>

#include <atlbase.h>
#include <Shlobj.h>

#include "utils.hpp"

class CIconStreams {
public:
    CIconStreams() = default;
    ~CIconStreams() = default;
public:
    struct IconInfo {
        HKEY hkey;
        tstring hkeystr;
        tstring subkey;
        tstring key;
        void* data;
        ULONG size;
        IconInfo() : data(nullptr), size(0) {}
        ~IconInfo() { if (data && size >0) { delete data; data = nullptr; size = 0; } }
    };
    typedef std::shared_ptr<IconInfo> ELEINFO;
    typedef std::vector<ELEINFO> MAPINFO;
public:
    void check() {
        static std::vector<tstring> hkeys = {
            _T("HKEY_CLASSES_ROOT\\Local Settings\\Software\\Microsoft\\Windows\\CurrentVersion\\TrayNotify"),
            _T("HKEY_CURRENT_USER\\Software\\Classes\\Local Settings\\Software\\Microsoft\\Windows\\CurrentVersion\\TrayNotify"),
        };
        static std::vector<tstring> keys = {
            _T("IconStreams"),
            _T("PastIconsStream"),
        };

        for (auto &hkey : hkeys) {
            for (auto &key : keys) {
                auto info = std::make_shared<IconInfo>();
                if (!this->parseHkey(hkey, info->hkey, info->subkey, info->hkeystr))
                    continue;
                info->key = key;
                if (!getHkey(info)) continue;
                map_.push_back(info);
            }
        }
    }

    MAPINFO& getMap() { return map_; }

    bool clear(const ELEINFO& ele) {
        /*for (auto p = (char*)ele->data; p - ele->data < ele->size; ++p) {
            
        }*/
        
        this->delHkey(ele->hkey, ele->subkey, ele->key);

        auto logstr = tformat(
            _T( "P:"), ele->hkeystr,
            _T(",S:"), ele->subkey,
            _T(",K:"), ele->key);
        tshow(logstr);
        tlog(logstr);
        
        return true;
    }

    void clearAll() {
        for (auto &ele : map_) {
            this->clear(ele);
        }
    }

    void reboot_explorer() {
        system("taskkill /im explorer.exe /f");

        Sleep(600);
        TCHAR szPath[255] = { 0 };
        if ((::SHGetFolderPath(NULL, CSIDL_WINDOWS, NULL, SHGFP_TYPE_CURRENT, szPath)) != S_OK)
            wsprintf(szPath, _T("c:\\windows"));
        ::ShellExecute(NULL, _T("open"), _T("explorer.exe"), 0, szPath, SW_NORMAL);
    }
protected:
    struct HKInfo {
        tstring key;
        HKEY hkey;
    };
    bool parseHkey(const tstring& hkey, HKEY& parent, tstring& subkey, tstring& hkeystr) {
        static std::vector<HKInfo> hkeys = {
            { _T("HKEY_CLASSES_ROOT"), HKEY_CLASSES_ROOT },
            { _T("HKEY_CURRENT_USER"), HKEY_CURRENT_USER },
            { _T("HKEY_LOCAL_MACHINE"), HKEY_LOCAL_MACHINE },
            { _T("HKEY_USERS"), HKEY_USERS },
            { _T("HKEY_PERFORMANCE_DATA"), HKEY_PERFORMANCE_DATA },
            { _T("HKEY_PERFORMANCE_TEXT"), HKEY_PERFORMANCE_TEXT },
            { _T("HKEY_PERFORMANCE_NLSTEXT"), HKEY_PERFORMANCE_NLSTEXT },
            { _T("HKEY_CURRENT_CONFIG"), HKEY_CURRENT_CONFIG },
            { _T("HKEY_DYN_DATA"), HKEY_DYN_DATA },
            { _T("HKEY_CURRENT_USER_LOCAL_SETTINGS"), HKEY_CURRENT_USER_LOCAL_SETTINGS }
        };

        auto pos = hkey.find(_T("\\"));
        if (pos == std::string::npos) return false;

        auto szp = hkey.substr(0, pos);
        for (auto it = hkeys.begin(); it != hkeys.end(); ++it) {
            if ((*it).key == szp) {
                hkeystr = szp;
                parent = (*it).hkey;
                subkey = hkey.substr(pos + 1, hkey.length() - pos);
                return true;
            }
        }

        return false;
    }
    bool getHkey(ELEINFO& info) {
        ATL::CRegKey reg;
        if (ERROR_SUCCESS != reg.Open(info->hkey, info->subkey.c_str(), KEY_READ))
            return false;

        if (ERROR_SUCCESS != reg.QueryBinaryValue(info->key.c_str(), nullptr, &info->size))
            return false;
        if (info->size <= 0) return false;

        info->data = new char[info->size];
        if (ERROR_SUCCESS != reg.QueryBinaryValue(info->key.c_str(), info->data, &info->size))
            return false;

        return true;        
    }
    bool setHkey(HKEY hkey, const tstring& subkey, const tstring& key, void* data, ULONG size) {
        ATL::CRegKey reg;
        if (ERROR_SUCCESS != reg.Open(hkey, subkey.c_str(), KEY_WRITE))
            return false;

        if (ERROR_SUCCESS != reg.SetBinaryValue(key.c_str(), data, size))
            return false;

        return true;
    }
    bool delHkey(HKEY hkey, const tstring& subkey, const tstring& key) {
        ATL::CRegKey reg;
        if (ERROR_SUCCESS != reg.Open(hkey, subkey.c_str(), KEY_WRITE))
            return false;

        if (ERROR_SUCCESS != reg.DeleteValue(key.c_str()))
            return false;

        return true;
    }
protected:
    MAPINFO map_;
};

#endif